from modelscope import AutoModelForCausalLM
from PIL import Image
import time
import os


# 使用绝对路径并确保格式正确
model_id = r"D:\moondream2"

# 检查模型目录和必要文件
if not os.path.exists(model_id):
    raise FileNotFoundError(f"模型目录不存在: {model_id}")

print(f"正在加载本地模型: {model_id}")

# 设置环境变量来禁用网络访问
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['MODELSCOPE_OFFLINE'] = '1'

# Load the model
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    device_map={"": "cuda"},  # 使用GPU加速
    local_files_only=True  # 强制只使用本地文件
)

# Load your image
image_path = "output_image.png"  # 替换为你的截图路径
image = Image.open(image_path)

# 定义系统提示和用户提示
system_prompt = '''
我正在玩一款名为"KitTechMatch3D"的游戏，这是一款经典的消除类游戏。游戏中有以下特点和规则:
        1.关卡的上方规定了要手机的物品种类和数量，在规定时间内收集完成后即可胜利。
        2.玩家需要使用鼠标点击来收集物品，收集的物品会出现在放置栏。
        3. 3个相同的物品在放置栏可以消除，如果放置栏满了，游戏会失败，当目标物体全部被消除后，游戏就会成功
        现在你是一个玩这种游戏的高手，你需要帮助玩家在尽可能短的时间内完成游戏目标的收集。根据游戏截图（截图里有一个个的格子，并带有编号），你需要给玩家建议点击哪个格子
'''


user_prompt = "我现在应该点击什么位置。只需要回答点击的位置(x, y)。"

time_start = time.time()

# 执行推理
result = model.query(image, system_prompt + "\n" + user_prompt)

print(f"predict cost time: {time.time() - time_start}s")
print(result["answer"])